#include "DTW.h"

//Class functions
Series* Series::createNewSeries(){  
  /*
    Heap allocate a series
   */
  Series *temp = new Series;
  temp->lengthSeries = 0;
  temp->classification = -1;
  temp->series = new double*[100];
  for (int j=0; j<100; j++){
    temp->series[j] = new double[13];
  }
  temp->mag = new double[100];
  temp->upperIndex = new int[100];
  temp->lowerIndex = new int[100];
  return temp;
}
void Series::freeMem(Series *ptr){
  /*
    Free the memory of a series
   */
  for (int i=0; i<100; i++){
    free(ptr->series[i]);
  }
  free(ptr->series);
  free(ptr->mag);
  free(ptr->upperIndex);
  free(ptr->lowerIndex);
  free(ptr);
}
void Series::addLine(string s, int line) {
  /*
    Add the coefficients in one frame to the series
   */
  stringstream ss;
  ss << s;
  for (int i=0; i<13; i++) {
    ss >> series[line][i];
  }
}
int Series::getMag(int i) {
  return mag[i];
}
int Series::setMag(int m, int i){
  /*
   Set the magnitude of the ith vector of the series to m
  */
  mag[i] = m;
}

int Series::getLength(){
  return lengthSeries;
}
void Series::setLength(int num){
  lengthSeries = num;
}
int * Series::getUpper(){
  return upperIndex;
}
int * Series::getLower(){
  return lowerIndex;
}
int Series::setUpper(int * upper) {
  for (int i=0; i<lengthSeries; i++)
    upperIndex[i] = upper[i];
}
int Series::setLower(int * lower) {
  for (int i=0; i<lengthSeries; i++)
    lowerIndex[i] = lower[i];
}
double ** Series::getSeries(){
  return series;
}
int Series::getClass(){
  return classification;
}
void Series::setClass(int c){
  classification = c;
}


Series **init(int datasetSize, int chunkSize, string filename, double sampleRate, int *finalSize){
  /*
    Read data from a file, truncate the unnecessary part if sampling the data.
   */
  char * fn = (char*) filename.c_str();
  double zero[13] = {0};
  Series **training = new Series *[datasetSize];
  for (int i=0; i<datasetSize; i++){
    training[i] = Series::createNewSeries();
  }
  ifstream fin;
  fin.open(fn);
  int i = 0, line = 0, mag[100] = {0};
  int count = 0, index = 0;
  while (i<datasetSize){
    string s;
    getline(fin, s);
    if (s.length() > 50){
      if (count < sampleRate * double(chunkSize)){
	training[index]->addLine(s, line);		
	training[index]->setMag(calDist(zero, training[index]->getSeries()[line]), line);
	line++;
      }
    }
    else{ 
      if (count < sampleRate * double(chunkSize)){
	training[index]->setClass(i / chunkSize); 
	training[index]->setLength(line);
	line = 0;	
	index += 1;
      }
      i++;
      count = (count + 1)%chunkSize;
    }
  }
  *finalSize = index;
  fin.close();
  return training; 
}

double calDist(double *pt1, double *pt2) {
  /*
    Calculate the Euclidean distance between two frames in two series
  */
  double dist = 0;
  for (int i=0; i<DIM; i++) 
    dist += (pt1[i] - pt2[i]) * (pt1[i] - pt2[i]);
  return dist;
}


//very first DTW
double dtw(double **trainBlock,int trainSize, double **testBlock, int testSize){\
  /*
    Only DTW
  */
  double DTW[100][100] = {0};
  DTW[0][0] = calDist(trainBlock[0], testBlock[0]);
  for (int i=1; i<trainSize; i++) 
    DTW[i][0] = calDist(trainBlock[i], testBlock[0]) + DTW[i-1][0];
  for (int i=1; i<testSize; i++)
    DTW[0][i] = calDist(trainBlock[0], testBlock[i]) + DTW[0][i-1];

  for(int i=1; i<trainSize; i++)
    for (int j=1;j<testSize; j++) {
      DTW[i][j] = min(DTW[i-1][j], min(DTW[i-1][j-1],DTW[i][j-1])) + calDist(trainBlock[i], testBlock[j]);
    }
  return DTW[trainSize-1][testSize-1];
  
}

double dtw(double **trainBlock,int trainSize, double **testBlock, int testSize, int window){
  /*
    DTW with Sakoe-Chiba band, take window as a parameter
  */
  int vis[100][100] = {0};
  double DTW[100][100] = {0};
  int visited[100][100] = {0};
  DTW[0][0] = calDist(trainBlock[0], testBlock[0]);
  DTW[trainSize-1][testSize-1] = BIG_NUMBER;
  visited[0][0] = 1;
  for (int i=1; i<=trainSize; i++){
    DTW[i][0] = calDist(trainBlock[i], testBlock[0]) + DTW[i-1][0];
    visited[i][0] = 1;
  }
  for (int i=1; i<=testSize; i++){
    DTW[0][i] = calDist(trainBlock[0], testBlock[i]) + DTW[0][i-1];
    visited[0][i] = 1;
  }
  for(int i=1; i<trainSize; i++) {
    int range = int(testSize / double(trainSize) * i);
    vis[i][range] = 1;
    for (int j=max(range-window,1);j<min(range+window,testSize); j++) {
	if (visited[i-1][j]==0) DTW[i-1][j] = BIG_NUMBER;
	if (visited[i-1][j-1]==0) DTW[i-1][j-1] = BIG_NUMBER; 
	if (visited[i][j-1]==0) DTW[i][j-1] = BIG_NUMBER;
	DTW[i][j] = min(DTW[i-1][j], min(DTW[i-1][j-1],DTW[i][j-1])) + calDist(trainBlock[i], testBlock[j]);
	visited[i][j] = 1;
      }
  }
  
  return DTW[trainSize-1][testSize-1];
}

void getUpperLowerIndex(Series *train, int window) {
  /*
    Set the upper/lower index for a series using the window value
  */
  double ** trainSeries = train->getSeries();
  int len = train->getLength();
  int upperIndex[100], lowerIndex[100];
  for (int i=0; i<len; i++) {
    int maxIndex=-1;
    int minIndex=-1;
    for (int j=max(0, i-window); j<min(len, i+window); j++) {
      if (maxIndex == -1 || train->getMag(j) > train->getMag(maxIndex)) 
	maxIndex = j;
      if (minIndex == -1 || train->getMag(j) < train->getMag(minIndex))
	minIndex = j;
    }
    upperIndex[i] = maxIndex;
    lowerIndex[i] = minIndex;
  }
  train->setUpper(upperIndex);
  train->setLower(lowerIndex);
}

double calcLowerBound(Series *train, Series *test, int window) {
  /*
    Calculate the lower bound of the distance between two series
  */
  double ** trainSeries = train->getSeries();
  double ** testSeries = test->getSeries();
  
  int len = min(train->getLength(), test->getLength());
  int * upperIndex = train->getUpper();
  int * lowerIndex = train->getLower();
  double lowerBound = 0;
  
  for (int i=0; i<len; i++) {
    double testMag = test->getMag(i);
    double upperMag = train->getMag(upperIndex[i]);
    double lowerMag = train->getMag(lowerIndex[i]);
    if (testMag > upperMag) 
      lowerBound += (testMag - upperMag) * (testMag - upperMag);
    else if (testMag < lowerMag) 
      lowerBound += (testMag-lowerMag) * (testMag - lowerMag);
  }
  return sqrt(lowerBound);
}

int kNN(Series **train, int trainSize, Series * entry, int option, int k, int window){
  /* 
     Perform k-nearest neighbor on the training set to classify an entry
  */
  double bestSoFar = BIG_NUMBER;
  int bestKindex[300] = {0};
  double minDist[300] = {0};
  int numK = 0;
  for (int i=0; i<trainSize; i++){
    double lowerBound;
    if (option == 2){
      lowerBound = calcLowerBound(train[i], entry, window);
      //uncomment the following lines if want to get a sense about which entries get cut
      //cout << "lb:" << lowerBound << " best:" << bestSoFar << endl;
      //if (lowerBound > bestSoFar)
      //cout << "cut!" << endl;
    }
    else double lowerBound = -1; // do not use lowerbound
    if (lowerBound < bestSoFar) {
      double dist;
      if (option > 0){
	if (train[i]->getLength() > entry->getLength())
	  dist = dtw(train[i]->getSeries(), train[i]->getLength(), entry->getSeries(), entry->getLength(), window);
	else
	  dist = dtw(entry->getSeries(), entry->getLength(), train[i]->getSeries(), train[i]->getLength(), window);
      }
      else
	dist = dtw(train[i]->getSeries(), train[i]->getLength(), entry->getSeries(), entry->getLength());
      if (numK<k){
	bestKindex[numK] = i;
	minDist[numK] = dist;
	numK ++;
      }
      else{
	double maxDist = dist;
	int maxIndex = -1;
	for (int j=0; j<k; j++)
	  if (maxDist < minDist[j]){
	    maxDist = minDist[j];
	    maxIndex = j;
	}
	if (maxIndex != -1){
	  minDist[maxIndex] = dist;
	  bestKindex[maxIndex] = i;
	  bestSoFar = -1;
	  for (int j=0; j<k; j++)
	    if (bestSoFar < minDist[j])
	      bestSoFar = minDist[j];
	}
      }
    }
  }
  int classify[10]={0};
  for (int i=0; i<k; i++){
    classify[train[bestKindex[i]]->getClass()]++;
  }
  int maxC = 0;
  for (int i=0; i<10; i++){
    if (classify[maxC]<classify[i])
      maxC = i;
  }
  return maxC;
}

void destroy(Series **setData, int setSize){
  /*
    Free a dataset
   */
  for (int i=0; i<setSize; i++){
    Series::freeMem(setData[i]);
  }
  free(setData);
}

void printUsage(){
  cout << "Dynamic Time Warping Program\n\n";
  cout << "usage: ./dtw <algorithm> <number-of-neighbors> <window-size> <sample-percentage>\n\n";
  cout << "A single <algorithm> must be specified:\n";
  cout << "\t--naive-DTW\n";
  cout << "\t--global-constraint\n";
  cout << "\t--lower-bound\n\n";
  cout << "A single <number-of-neighbors> must be specified. It must be an integer\n";
  cout << "chosen from 1 to 50.\n\n";
  cout << "If --global-constraint or --lower-bound is specified,\n";
  cout << "<window-size> must be specified. It must be an integer chosen from 1 to 50.\n\n";
  cout << "A single <sample-percentage> must be specified. It must satisfy that\n0 < <sample-percentage> <= 1.\n\n";
  exit(0);
}

void parseCommand(int argc, char *argv[], int *option, int *window, int *k, double *sampleRate) {
  /*
    Parse the argument from command line
  */
  if (argc!=4 && argc!=5)
    printUsage();
  if (argc==4) {
    if (strncmp(argv[1],"--naive",50) || atoi(argv[2])<1 || atoi(argv[2])>50 || 
	atof(argv[3])<=0 ||atof(argv[3])>1)
      printUsage();
    *option = 0;
    *window = 0;
    *k = atoi(argv[2]);
    *sampleRate = atof(argv[3]);
  }
  if (argc==5) {
    if (strncmp(argv[1], "--global-constraint",50) && strncmp(argv[1],"--lower-bound",50))
      printUsage();
    if (atoi(argv[2])<1 || atoi(argv[2])>50 || atoi(argv[3])<1 || atoi(argv[3])>50 
	||atof(argv[4])<=0 || atof(argv[4])>1)
      printUsage();
    if (!strncmp(argv[1], "--global-constraint",50))
      *option = 1;
    else
      *option = 2;
    *k = atoi(argv[2]);
    *window = atoi(argv[3]);
    *sampleRate = atof(argv[4]);
  }
}


double getAccuracy(Series **train, int trainingSize, Series **test, int testSize,int option, int k, int window){
  /*
    Return the accuracy when classify a set against the training set
  */
  int count = 0;
  for (int i=0; i<testSize; i++){
    int predict = 0;
    predict = kNN(train, trainingSize, test[i], option, k, window);
    //cout << i << " " << test[i]->getClass() << " " << predict << endl;
    if (predict == test[i]->getClass()){
      count += 1;
    }
  }
  return count/double(testSize);
}

int main(int argc, char* argv[]){
  int option,window,k;
  double sampleRate;
  
  parseCommand(argc, argv, &option, &window, &k, &sampleRate);
  
  int trainSize=0, testSize=0;
  Series **training = init(6600, 660, "Train_Arabic_Digit.txt", sampleRate, &trainSize);
  if (option==2)
    for (int i=0; i<6600; i++)
      getUpperLowerIndex(training[i], window);
  Series **test = init(2200, 220, "Test_Arabic_Digit.txt", sampleRate, &testSize);
  
  cout << "finish reading data" << endl;

  cout << "Test Set Accuracy for k="<<k << "\t" << getAccuracy(training, trainSize, test, testSize, option, k, window) << endl;
  destroy(training,trainSize);
  destroy(test,testSize);
  
  return 0;
}
